package net.xt.testthread.utils;

import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 自定义性能测试工具类
 */
@Slf4j
public class LoadRunner {

    /**
     * 运行自定义性能测试
     *
     * @param requests    请求数
     * @param concurrency 并发数
     * @param runnable    自定义执行的任务方法
     * @return LoadRunnerResult
     * @throws InterruptedException
     */
    public static LoadRunnerResult run(int requests, int concurrency, Runnable runnable)
            throws InterruptedException {
        log.info("开始性能测试......");

        // 创建线程池并将所有核心线程池都准备好，core=max，直接设置到最大
        ThreadPoolExecutor threadPool = new ThreadPoolExecutor(
                concurrency,
                concurrency,
                0L,
                TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<>());
        // 启动所有核心线程
        threadPool.prestartAllCoreThreads();

        // 创建线程计数器，默认大小=请求数量
        CountDownLatch countDownLatch = new CountDownLatch(requests);

        // 定义成功请求数，最快耗时，最慢耗时三个原子变量，防止多线程并发操作变量值错误
        AtomicInteger successRequests = new AtomicInteger(0);
        AtomicLong fastestCostTime = new AtomicLong(Integer.MAX_VALUE);
        AtomicLong slowestCostTime = new AtomicLong(Integer.MIN_VALUE);

        // 记录开始时间
        long startTime = System.currentTimeMillis();
        try {
            for (int i = 0; i < requests; i++) {
                threadPool.execute(() -> {
                    try {
                        // 记录每个请求任务的开始时间
                        long requestStartTime = System.currentTimeMillis();
                        // 任务执行
                        runnable.run();

                        // 每个请求任务执行耗时
                        long costTime = System.currentTimeMillis() - requestStartTime;
                        // 请求最快耗时
                        setFastestCostTime(fastestCostTime, costTime);
                        // 请求最慢耗时
                        setSlowestCostTimeCostTime(slowestCostTime, costTime);
                        // 成功请求数+1
                        successRequests.incrementAndGet();
                    } catch (Exception e) {
                        log.error(e.getMessage());
                    } finally {
                        // 计数器-1
                        countDownLatch.countDown();
                    }
                });
            }

            // 阻塞当前线程，等所有线程结束后，再继续向下走
            countDownLatch.await();
        } catch (Exception e) {
            log.error(e.getMessage());
        } finally {
            // 关闭线程池
            threadPool.shutdown();
        }

        // 记录结束时间
        long endTime = System.currentTimeMillis();
        log.info("性能测试结束，总耗时(ms):{}", endTime - startTime);

        // 组装最后的结果并返回
        LoadRunnerResult loadRunnerResult = new LoadRunnerResult();
        loadRunnerResult.setRequests(requests);
        loadRunnerResult.setConcurrency(concurrency);
        loadRunnerResult.setSuccessRequests(successRequests.get());
        loadRunnerResult.setFailRequests(requests - loadRunnerResult.getSuccessRequests());
        loadRunnerResult.setTotalCostTime(endTime - startTime);
        loadRunnerResult.setRequestsPerSecond((float) requests * 1000f / (float) (loadRunnerResult.getTotalCostTime()));
        loadRunnerResult.setTimePerRequest((float) loadRunnerResult.getTotalCostTime() / requests);
        loadRunnerResult.setFastestCostTime(fastestCostTime.get());
        loadRunnerResult.setSlowestCostTime(slowestCostTime.get());
        return loadRunnerResult;
    }


    private static void setFastestCostTime(AtomicLong fastestCostTime, long costTime) {
        while (true) {
            long fsCostTime = fastestCostTime.get();
            if (fsCostTime < costTime) {
                break;
            }

            if (fastestCostTime.compareAndSet(fsCostTime, costTime)) {
                break;
            }
        }
    }

    private static void setSlowestCostTimeCostTime(AtomicLong slowestCostTime, long costTime) {
        while (true) {
            long slCostTime = slowestCostTime.get();
            if (slCostTime > costTime) {
                break;
            }

            if (slowestCostTime.compareAndSet(slCostTime, costTime)) {
                break;
            }
        }
    }
}
